import pickle
import jax
import matplotlib.pyplot as plt
import jax.numpy as jnp
import tensorflow_probability.substrates.jax as tfp
from scipy.stats import gaussian_kde
import plotly.express as px
import pandas as pd
import pickle
tfd = tfp.distributions
import plotly
from laplax import ADLaplace
plotly.offline.init_notebook_mode()
2022-06-22 17:28:04.702758: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
def plot_all(varient=''):
all_pdfs = []
all_labels = []
varient= str(varient)
x = jnp.linspace(0,6,10000)
with open('./results_data/linear_regression_Ajax'+varient,'rb') as f:
variational = pickle.load(f)
params = variational.get_params()
loc_m, scale = jax.tree_leaves(variational.transform_dist(params['theta']))
scale = jnp.dot(scale, scale.T)
for i in range(2):
y = tfd.Normal(loc = loc_m[i],scale = jnp.sqrt(scale[i][i])).prob(x)
all_pdfs.append(y)
all_labels.append('Ajax VI theta0')
all_labels.append('Ajax VI theta1')
with open('./results_data/linear_regression_laplace'+varient,'rb') as f:
laplace = pickle.load(f)
loc_m = laplace['mean']
std = jnp.sqrt(jnp.diag(laplace['cov']))
for i in range(2):
y = tfd.Normal(loc = loc_m[i],scale = std[i]).prob(x)
all_pdfs.append(y)
all_labels.append('Laplace approximation theta0')
all_labels.append('Laplace approximation theta1')
with open('./results_data/MCMC_Blackjax'+varient,'rb') as f:
black_samples = pickle.load(f)
for i in range(2):
kde_black = gaussian_kde(black_samples.position['theta'][:,i])
pdf_black = kde_black(x)
all_pdfs.append(pdf_black)
all_labels.append('Blackjax rmh theta0')
all_labels.append( 'Blackjax rmh theta1')
all_pdfs = jnp.array(all_pdfs).reshape((-1))
no_estimates = len(all_labels)
all_labels_repeated = [item for item in all_labels for i in range(x.shape[0])]
x_repeated = jnp.tile(x,no_estimates)
to_df = {
"theta":x_repeated,
"PDF":all_pdfs,
"label": all_labels_repeated
}
df = pd.DataFrame(to_df)
fig = px.line(to_df,"theta","PDF",color="label",title=f"Linear regression posterior")
fig.show()
plot_all()
!jupyter nbconvert --to HTML linear_regression_results.ipynb